from dataclasses import dataclass
from typing import List, Optional, Dict, Any
import bencodepy
import time

from src.routing.routing_table import Node

@dataclass
class DHTMessage:
    """DHT基础消息类
    
    所有DHT消息的基类，实现基本的序列化和反序列化功能。
    
    Attributes:
        transaction_id (bytes): 事务ID，用于匹配请求和响应
        node_id (bytes): 发送节点的ID，20字节
        message_type (str): 消息类型，'q'=查询, 'r'=响应, 'e'=错误
    """
    transaction_id: bytes
    node_id: bytes
    message_type: str
    
    def to_bytes(self) -> bytes:
        """将消息转换为bencode编码的字节序列"""
        data = {
            't': self.transaction_id,
            'y': self.message_type,
            'v': '1.0'
        }
        return bencodepy.encode(data)
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'DHTMessage':
        """从bencode编码的字节创建消息对象"""
        decoded = bencodepy.decode(data)
        return cls(
            transaction_id=decoded[b't'],
            node_id=decoded[b'a'][b'id'] if b'a' in decoded else decoded[b'r'][b'id'],
            message_type=decoded[b'y'].decode()
        )

@dataclass
class QueryMessage(DHTMessage):
    """查询消息基类
    
    所有查询类型消息的基类，如 ping, find_node 等。
    
    Attributes:
        query_name (str): 查询方法名称
    """
    query_name: str
    
    def __init__(self, transaction_id: bytes, node_id: bytes, query_name: str):
        """初始化查询消息
        
        Args:
            transaction_id: 事务ID
            node_id: 节点ID
            query_name: 查询方法名称
        """
        super().__init__(
            transaction_id=transaction_id,
            node_id=node_id,
            message_type='q'
        )
        self.query_name = query_name
    
    def add_args(self,args:Dict[str,Any]) -> None:
        """添加参数"""
        pass
    
    def to_bytes(self) -> bytes:
        """序列化查询消息"""
        data = {
            't': self.transaction_id,
            'y': 'q',
            'q': self.query_name,
            'a': {
                'id': self.node_id
            }
        }
        self.add_args(data)
        return bencodepy.encode(data)

@dataclass
class ResponseMessage(DHTMessage):
    """响应消息基类
    
    所有响应类型消息的基类。
    
    Attributes:
        response_data (Dict[bytes, Any]): 响应数据字典
    """
    response_data: Dict[bytes, Any] = None
    
    def __init__(self, transaction_id: bytes, node_id: bytes):
        """初始化响应消息
        
        Args:
            transaction_id: 事务ID
            node_id: 节点ID
        """
        super().__init__(
            transaction_id=transaction_id,
            node_id=node_id,
            message_type='r'
        )
        self.response_data = {}
    
    def to_bytes(self) -> bytes:
        """序列化响应消息"""
        data = {
            't': self.transaction_id,
            'y': 'r',
            'r': {
                'id': self.node_id,
                **self.response_data
            }
        }
        return bencodepy.encode(data)

    @staticmethod
    def from_bytes(data: bytes) -> 'ResponseMessage':
        """从字节数据反序列化响应消息
        
        Args:
            data: 待解析的字节数据
            
        Returns:
            ResponseMessage: 解析后的响应消息对象
            
        Raises:
            ValueError: 数据格式不正确时抛出
        """
        try:
            decoded = bencodepy.decode(data)
            if not isinstance(decoded, dict):
                raise ValueError("Invalid message format")
            
            if decoded.get(b'y') != b'r':
                raise ValueError("Not a response message")
            
            transaction_id = decoded.get(b't')
            if not transaction_id:
                raise ValueError("Missing transaction ID")
            
            response_data = decoded.get(b'r', {})
            if not isinstance(response_data, dict):
                raise ValueError("Invalid response data format")
            
            node_id = response_data.get(b'id')
            if not node_id:
                raise ValueError("Missing node ID")
            
            # 创建响应消息实例
            message = ResponseMessage(transaction_id, node_id)
            # 移除node_id，剩余数据存入response_data
            response_data.pop(b'id', None)
            message.response_data = response_data
            
            return message
            
        except bencodepy.DecodingError as e:
            raise ValueError(f"Failed to decode message: {str(e)}")

@dataclass
class PingMessage(QueryMessage):
    """PING查询消息"""
    def __init__(self, transaction_id: bytes, node_id: bytes):
        super().__init__(
            transaction_id=transaction_id,
            node_id=node_id,
            query_name='ping'
        )

@dataclass
class PongMessage(ResponseMessage):
    """PING响应消息"""
    pass

@dataclass
class FindNodeMessage(QueryMessage):
    """查找节点查询消息
    
    Attributes:
        target (bytes): 目标节点ID，20字节
    """
    target: bytes
    
    def __init__(self, transaction_id: bytes, node_id: bytes, target: bytes):
        super().__init__(
            transaction_id=transaction_id,
            node_id=node_id,
            query_name='find_node'
        )
        self.target = target
    
    def to_bytes(self) -> bytes:
        """序列化find_node消息"""
        data = super().to_bytes()
        decoded = bencodepy.decode(data)
        decoded[b'a'][b'target'] = self.target
        return bencodepy.encode(decoded)

@dataclass
class NodesMessage(ResponseMessage):
    """Nodes响应消息
    
    用于返回最接近目标的K个节点信息。
    
    消息格式:
    {
        "t": "aa",                     # 事务ID
        "y": "r",                      # 消息类型(响应)
        "r": {
            "id": "abcdef...",         # 20字节的节点ID
            "nodes": "def456..."       # 最接近的nodes信息
        }
    }
    """
    
    def __init__(self, transaction_id: bytes, node_id: bytes, nodes: bytes):
        super().__init__(transaction_id, node_id)
        self.nodes = nodes
        self.response_data = {
            b'nodes': nodes
        }

    def parse_nodes(self) -> List[Node]:
        """解析节点信息
        
        Returns:
            List[Node]: 解析出的节点列表
        """
        nodes = []
        nodes_data = self.nodes
        
        # 每个节点占26字节: 20字节node_id + 4字节ip + 2字节port
        for i in range(0, len(nodes_data), 26):
            node_data = nodes_data[i:i+26]
            if len(node_data) == 26:
                node_id = node_data[:20]
                ip = '.'.join([str(b) for b in node_data[20:24]])
                port = int.from_bytes(node_data[24:26], 'big')
                
                nodes.append(Node(
                    node_id=node_id,
                    ip=ip,
                    port=port,
                    last_seen=time.time()
                ))
        
        return nodes

    @classmethod
    def from_bytes(cls, data: bytes) -> 'NodesMessage':
        """从bencode编码的字节创建Nodes响应消息
        
        Args:
            data: bencode编码的响应数据
            
        Returns:
            NodesMessage: 解析后的响应消息
            
        Raises:
            ValueError: 数据格式不正确时抛出
        """
        try:
            decoded = bencodepy.decode(data)
            if not isinstance(decoded, dict):
                raise ValueError("Invalid message format")
            
            if decoded.get(b'y') != b'r':
                raise ValueError("Not a response message")
            
            transaction_id = decoded.get(b't')
            if not transaction_id:
                raise ValueError("Missing transaction ID")
            
            response_data = decoded.get(b'r', {})
            if not isinstance(response_data, dict):
                raise ValueError("Invalid response data format")
            
            node_id = response_data.get(b'id')
            if not node_id:
                raise ValueError("Missing node ID")
            
            nodes = response_data.get(b'nodes', b'')
            
            return cls(
                transaction_id=transaction_id,
                node_id=node_id,
                nodes=nodes
            )
            
        except bencodepy.DecodingError as e:
            raise ValueError(f"Failed to decode message: {str(e)}")

@dataclass
class GetPeersMessage(QueryMessage):
    """获取节点对等节点查询消息
    
    用于在DHT网络中查询指定info_hash的peers信息。
    
    消息格式:
    {
        "t": "aa",                     # 事务ID
        "y": "q",                      # 消息类型(查询)
        "q": "get_peers",              # 查询方法
        "a": {
            "id": "abcdef...",         # 20字节的节点ID
            "info_hash": "mnopqr..."   # 20字节的目标info_hash
        }
    }
    
    Attributes:
        info_hash (bytes): 要查询的目标info_hash，20字节
    """
    info_hash: bytes
    
    def __init__(self, transaction_id: bytes, node_id: bytes, info_hash: bytes):
        """初始化get_peers查询消息
        
        Args:
            transaction_id (bytes): 事务ID，用于匹配请求和响应
            node_id (bytes): 发送节点的ID，20字节
            info_hash (bytes): 要查询的目标info_hash，20字节
            
        Example:
            >>> msg = GetPeersMessage(
            ...     transaction_id=b'aa',
            ...     node_id=b'abcdefghij0123456789',
            ...     info_hash=b'mnopqrstuvwxyz123456'
            ... )
        """
        super().__init__(transaction_id, node_id, 'get_peers')
        self.info_hash = info_hash
    
    def add_args(self,args:Dict[str,Any]) -> None:
        """添加参数"""
        args['a']['info_hash'] = self.info_hash

@dataclass
class GetPeersResponseMessage(ResponseMessage):
    """GetPeers响应消息
    
    用于解析和构造get_peers查询的响应，可能包含peers列表或最接近的nodes。
    
    消息格式:
    1. 包含peers:
    {
        "t": "aa",
        "y": "r",
        "r": {
            "id": "abcdef...",         # 20字节的节点ID
            "token": "aoeusnth",       # token用于announce_peer
            "values": ["peer1", ...]   # peers列表
        }
    }
    
    2. 包含nodes:
    {
        "t": "aa",
        "y": "r",
        "r": {
            "id": "abcdef...",         # 20字节的节点ID
            "token": "aoeusnth",       # token用于announce_peer
            "nodes": "def456..."       # 最接近的nodes信息
        }
    }
    
    Attributes:
        transaction_id (bytes): 事务ID
        node_id (bytes): 响应节点的ID
        response_data (Dict[bytes, Any]): 响应数据
    """
    
    def __init__(self, transaction_id: bytes, node_id: bytes, token: bytes,
                 peers: Optional[List[bytes]] = None, nodes: Optional[bytes] = None):
        """初始化GetPeers响应消息
        
        Args:
            transaction_id: 事务ID
            node_id: 响应节点的ID
            token: 用于announce_peer的token
            peers: peers列表，每个peer为6字节(IP+端口)
            nodes: 最接近的nodes信息，每个node为26字节
            
        Example:
            >>> # 构造包含peers的响应
            >>> msg = GetPeersResponseMessage(
            ...     transaction_id=b'aa',
            ...     node_id=b'abcdefghij0123456789',
            ...     token=b'aoeusnth',
            ...     peers=[b'axje.u', b'idhtnm']
            ... )
            
            >>> # 构造包含nodes的响应
            >>> msg = GetPeersResponseMessage(
            ...     transaction_id=b'aa',
            ...     node_id=b'abcdefghij0123456789',
            ...     token=b'aoeusnth',
            ...     nodes=b'def456...'
            ... )
        """
        super().__init__(transaction_id, node_id)
        
        # 构建响应数据
        self.response_data = {
            b'id': node_id,
            b'token': token
        }
        
        if peers is not None:
            self.response_data[b'values'] = peers
        if nodes is not None:
            self.response_data[b'nodes'] = nodes
    
    @property
    def token(self) -> bytes:
        """获取响应中的token"""
        return self.response_data[b'token']
    
    @property
    def peers(self) -> Optional[List[bytes]]:
        """获取响应中的peers列表"""
        return self.response_data.get(b'values')
    
    @property
    def nodes(self) -> Optional[bytes]:
        """获取响应中的nodes信息"""
        return self.response_data.get(b'nodes')
    
    @classmethod
    def from_bytes(cls, data: bytes) -> 'GetPeersResponseMessage':
        """从bencode编码的字节创建GetPeers响应消息
        
        Args:
            data: bencode编码的响应数据
            
        Returns:
            GetPeersResponseMessage: 解析后的响应消息
            
        Raises:
            BencodeDecodeError: bencode解码失败
            KeyError: 缺少必要字段
            
        Example:
            >>> data = b'd1:rd2:id20:abcdefghij01234567895:token8:aoeusnth6:valuesl6:axje.u6:idhtnmee1:t2:aa1:y1:re'
            >>> resp = GetPeersResponseMessage.from_bytes(data)
        """
        decoded = bencodepy.decode(data)
        if b'r' not in decoded or b't' not in decoded:
            raise KeyError("Invalid response format")
            
        r_data = decoded[b'r']
        if b'token' not in r_data:
            raise KeyError("Missing token in get_peers response")
            
        # 从解析的数据构造响应对象
        return cls(
            transaction_id=decoded[b't'],
            node_id=r_data[b'id'],
            token=r_data[b'token'],
            peers=r_data.get(b'values'),
            nodes=r_data.get(b'nodes')
        )
    
    @classmethod
    def create_with_peers(cls, transaction_id: bytes, node_id: bytes, 
                         token: bytes, peers: List[bytes]) -> 'GetPeersResponseMessage':
        """创建包含peers的响应消息
        
        Args:
            transaction_id: 事务ID
            node_id: 节点ID
            token: 用于announce_peer的token
            peers: peers列表
        """
        return cls(transaction_id, node_id, token, peers=peers)
    
    @classmethod
    def create_with_nodes(cls, transaction_id: bytes, node_id: bytes,
                         token: bytes, nodes: bytes) -> 'GetPeersResponseMessage':
        """创建包含nodes的响应消息
        
        Args:
            transaction_id: 事务ID
            node_id: 节点ID
            token: 用于announce_peer的token
            nodes: 最接近的nodes信息
        """
        return cls(transaction_id, node_id, token, nodes=nodes)

@dataclass
class AnnouncePeerMessage(QueryMessage):
    """Announce_peer查询消息
    
    用于向DHT网络宣告自己正在下载某个资源。
    
    消息格式:
    {
        "t": "aa",                     # 事务ID
        "y": "q",                      # 消息类型(查询)
        "q": "announce_peer",          # 查询方法
        "a": {
            "id": "abcdef...",         # 20字节的节点ID
            "info_hash": "mnopqr...",  # 20字节的资源info_hash
            "port": 6881,              # 下载端口
            "token": "aoeusnth",       # 之前get_peers返回的token
            "implied_port": 0          # 可选，1表示使用数据包的源端口
        }
    }
    """
    info_hash: bytes
    port: int
    token: bytes
    implied_port: int = 0
    
    def __init__(self, transaction_id: bytes, node_id: bytes, info_hash: bytes, 
                 port: int, token: bytes, implied_port: int = 0):
        """初始化announce_peer查询消息
        
        Args:
            transaction_id: 事务ID
            node_id: 发送节点的ID
            info_hash: 资源的info_hash
            port: 下载端口
            token: get_peers响应中返回的token
            implied_port: 是否使用数据包的源端口
        """
        super().__init__(transaction_id, node_id, 'announce_peer')
        self.info_hash = info_hash
        self.port = port
        self.token = token
        self.implied_port = implied_port

    def add_args(self, args: Dict[str, Any]) -> None:
        """添加查询参数"""
        args['a'].update({
            'info_hash': self.info_hash,
            'port': self.port,
            'token': self.token,
            'implied_port': self.implied_port
        })

    @classmethod
    def from_bytes(cls, data: bytes) -> 'AnnouncePeerMessage':
        """从bencode编码的字节创建announce_peer查询消息
        
        Args:
            data: bencode编码的查询数据
            
        Returns:
            AnnouncePeerMessage: 解析后的查询消息
            
        Raises:
            ValueError: 数据格式不正确时抛出
        """
        try:
            decoded = bencodepy.decode(data)
            if not isinstance(decoded, dict):
                raise ValueError("Invalid message format")
            
            if decoded.get(b'y') != b'q' or decoded.get(b'q') != b'announce_peer':
                raise ValueError("Not an announce_peer query")
            
            args = decoded.get(b'a', {})
            if not isinstance(args, dict):
                raise ValueError("Invalid arguments format")
            
            return cls(
                transaction_id=decoded.get(b't'),
                node_id=args.get(b'id'),
                info_hash=args.get(b'info_hash'),
                port=args.get(b'port', 0),
                token=args.get(b'token'),
                implied_port=args.get(b'implied_port', 0)
            )
            
        except bencodepy.DecodingError as e:
            raise ValueError(f"Failed to decode message: {str(e)}")

@dataclass
class AnnouncePeerResponseMessage(ResponseMessage):
    """Announce_peer响应消息
    
    消息格式:
    {
        "t": "aa",                     # 事务ID
        "y": "r",                      # 消息类型(响应)
        "r": {
            "id": "abcdef..."          # 20字节的节点ID
        }
    }
    """
    
    def __init__(self, transaction_id: bytes, node_id: bytes):
        """初始化announce_peer响应消息
        
        Args:
            transaction_id: 事务ID
            node_id: 响应节点的ID
        """
        super().__init__(transaction_id, node_id)

    @classmethod
    def from_bytes(cls, data: bytes) -> 'AnnouncePeerResponseMessage':
        """从bencode编码的字节创建announce_peer响应消息
        
        Args:
            data: bencode编码的响应数据
            
        Returns:
            AnnouncePeerResponseMessage: 解析后的响应消息
            
        Raises:
            ValueError: 数据格式不正确时抛出
        """
        try:
            decoded = bencodepy.decode(data)
            if not isinstance(decoded, dict):
                raise ValueError("Invalid message format")
            
            if decoded.get(b'y') != b'r':
                raise ValueError("Not a response message")
            
            response = decoded.get(b'r', {})
            if not isinstance(response, dict):
                raise ValueError("Invalid response format")
            
            return cls(
                transaction_id=decoded.get(b't'),
                node_id=response.get(b'id')
            )
            
        except bencodepy.DecodingError as e:
            raise ValueError(f"Failed to decode message: {str(e)}")
